#include "nfa.h"
#include <cassert>

namespace regex
{
Anchor Input::get_anchor() const
{
    if (m_string.empty())
    {
        return static_cast<Anchor>(kBegin | kEnd);
    }
    
    if (m_index == 0)
    {
        return kBegin;
    }
    
    if (m_index + 1 >= m_string.size())
    {
        return kEnd;
    }
    
    return kNone;
}
    
bool Trigger::match_only_anchor(const Input &input) const
{
    return (m_anchor & input.get_anchor()) == m_anchor;
}

bool Trigger::match(const Input &input) const
{
    assert(match_only_anchor(input));

    if (input.is_epsilon())
    {
        return is_epsilon();
    }
    
    return std::binary_search(
        m_string.begin(),
        m_string.end(),
        input.get_char()
        );
}

Nfa::Nfa()
    : m_curr(0),
      m_dirty(false)
{
}

uint Nfa::get_start_state() const
{
    return 0;
}

bool Nfa::get_next_states(uint curr, const Input &input, std::vector<uint> &next) const
{
    if (m_dirty)
    {
        // Caller forgot to call finalize(), so force to call it here.
        const_cast<Nfa *>(this)->finalize();
    }

    uint start = 0;
    uint end = 0;
    find_match_transitions(curr, start, end);

    bool ret = false;
    for (size_t i = start; i < end; ++i)
    {
        const Transition &trans = m_transitions[i];
        const Trigger &trigger = trans.m_trigger;

        // Anchor must be matched if exists
        if (!trigger.match_only_anchor(input))
        {
            continue;
        }
        
        // Try to match them first
        if (trigger.match(input))
        {
            next.push_back(trans.m_next);
            get_next_states(trans.m_next, input.epsilon(), next);
            ret = true;
            continue;
        }
        
        // Still continue the transition if the trigger is epsilon
        if (trigger.is_epsilon())
        {
            ret = get_next_states(trans.m_next, input, next) || ret;
            continue;
        }
    }
    return ret;
}

bool Nfa::is_accept(uint u) const
{
    if (m_dirty)
    {
        // Caller forgot to call finalize(), so force to call it here.
        const_cast<Nfa *>(this)->finalize();
    }

    return std::binary_search(
        m_accepts.begin(), 
        m_accepts.end(), 
        u
        );
}

uint Nfa::add_state()
{
    return ++m_curr;
}

void Nfa::add_transition(uint curr, uint next, const Trigger &trigger)
{
    assert(!trigger.is_epsilon() || curr != next);

    Transition trans(curr, next, trigger);
    m_transitions.push_back(trans);
    m_dirty = true;
}

void Nfa::add_epsilon_transition(uint curr, uint next)
{
    if (curr != next)
    {
        Trigger epsilon("", kNone);
        add_transition(curr, next, epsilon);
    }
}

void Nfa::set_accept(uint u)
{
    m_accepts.push_back(u);
    m_dirty = true;
}

void Nfa::finalize()
{
    std::sort(m_accepts.begin(), m_accepts.end());
    std::sort(m_transitions.begin(), m_transitions.end(), less_transition);
    m_dirty = false;
}

void Nfa::clear()
{
    m_curr = 0;
    m_dirty = false;
    m_transitions.clear();
    m_accepts.clear();
}

bool Nfa::less_transition(const Transition &left, const Transition &right)
{
    return left.m_curr < right.m_curr;
}

bool Nfa::get_transitions(uint curr, std::vector<Transition> &trans) const
{
    uint start = 0;
    uint end = 0;
    find_match_transitions(curr, start, end);

    for (size_t i = start; i < end; ++i)
    {
        trans.push_back(m_transitions[i]);
    }
    
    return !trans.empty();
}

bool Nfa::find_match_transitions(uint curr, uint &start, uint &end) const
{
    // Binary search the matched transitions
    start = 0;
    end = m_transitions.size();
    while (start < end)
    {
        size_t mid = (start + end) / 2;
        if (m_transitions[mid].m_curr < curr)
        {
            start = mid + 1;
        }
        else
        {
            end = mid;
        }
    }
    while (end < m_transitions.size() && m_transitions[end].m_curr == curr)
    {
        ++end;
    }
    
    return start < end;
}

bool Nfa::has_accept(const std::vector<uint> &states) const
{
    for (size_t i = 0; i < states.size(); ++i)
    {
        if (is_accept(states[i]))
        {
            return true;
        }
    }
    return false;
}

bool Nfa::remove_accept(std::vector<uint> &states) const
{
    bool found = false;
    for (size_t i = states.size(); i > 0; --i)
    {
        if (is_accept(states[i - 1]))
        {
            states.erase(states.begin() + i - 1);
            found = true;
        }
    }
    return found;
}

bool NfaMatcher::match(const Nfa &nfa, const std::string &str, MatchResult &mr)
{
    std::vector<uint> curr;
    std::vector<uint> next;

    mr.m_original = &str;
    mr.m_start = 0;
    mr.m_end = 0;
    bool ever_matched = false;

    if (nfa.is_accept(nfa.get_start_state()))
    {
        ever_matched = true;
    }

    Input input(str, 0);
    if (nfa.get_next_states(nfa.get_start_state(), input.epsilon(), next))
    {
        ever_matched = nfa.has_accept(next);
        next.clear();
    }

    curr.push_back(nfa.get_start_state());
    for (size_t i = 0; i <= str.size(); ++i)
    {
        Input input(str, i);
        for (size_t j = 0; j < curr.size(); ++j)
        {
            nfa.get_next_states(curr[j], input, next);
        }

        // Uniq next states
        std::sort(next.begin(), next.end());
        std::vector<uint>::iterator new_end = std::unique(next.begin(), next.end());
        next.erase(new_end, next.end());

        // Swap next states to curr for next loop.
        next.swap(curr);
        next.clear();

        if (nfa.remove_accept(curr))
        {
            mr.m_end = i + 1;
            ever_matched = true;
        }

        if (curr.empty())
        {
            break;
        }
    }
    
    return ever_matched;
}

bool NfaBuilder::build(const std::string &str, Nfa &nfa)
{
    nfa.clear();

    std::vector<uint> group_start;
    std::vector<uint> group_end;
    std::vector<size_t> group_end_index;

    uint prev_start = nfa.get_start_state();
    uint prev_end = nfa.get_start_state();
    bool permit_repeat = false;

    group_start.push_back(prev_end);

    for (uint i = 0; i < str.size(); ++i)
    {
        char ch = str[i];
        switch (ch)
        {
            case '(':
                group_start.push_back(prev_end);
                group_end_index.push_back(group_end.size());
                permit_repeat = false;
                break;
            case ')':
                if (group_end_index.empty())
                {
                    /// Todo: log error
                    return false;
                }
                for (size_t j = group_end_index.back(); j < group_end.size(); ++j)
                {
                    nfa.add_epsilon_transition(group_end[j], prev_end);
                }
                prev_start = group_start.back();
                group_start.pop_back();
                group_end.resize(group_end_index.back());
                group_end_index.pop_back();
                permit_repeat = true;
                break;
            case '|':
                group_end.push_back(prev_end);
                prev_end = group_start.back();
                permit_repeat = false;
                break;
            case '[':
                {
                    std::string chars;
                    if (!parse_char_set(str, ++i, chars))
                    {
                        // TODO: Log error
                        return false;
                    }
                    Trigger trigger(chars, kNone);
                    uint new_state = nfa.add_state();
                    nfa.add_transition(prev_end, new_state, trigger);
                    prev_start = prev_end;
                    prev_end = new_state;
                    permit_repeat = true;
                }
                break;
            case '^':
                {
                    Trigger trigger("", kBegin);
                    uint new_state = nfa.add_state();
                    nfa.add_transition(prev_end, new_state, trigger);
                    prev_start = prev_end;
                    prev_end = new_state;
                    permit_repeat = false;
                }
                break;
            case '$':
                {
                    Trigger trigger("", kEnd);
                    uint new_state = nfa.add_state();
                    nfa.add_transition(prev_end, new_state, trigger);
                    prev_start = prev_end;
                    prev_end = new_state;
                    permit_repeat = false;
                }
                break;
            case '+':
                if (!permit_repeat)
                {
                    // TODO: Log error
                    return false;
                }
                nfa.add_epsilon_transition(prev_end, prev_start);
                permit_repeat = false;
                break;
            case '*':
                if (!permit_repeat)
                {
                    // TODO: Log error
                    return false;
                }
                nfa.add_epsilon_transition(prev_end, prev_start);
                prev_end = prev_start;
                permit_repeat = false;
                break;
            case '?':
                if (!permit_repeat)
                {
                    // TODO: Log error
                    return false;
                }
                nfa.add_epsilon_transition(prev_start, prev_end);
                permit_repeat = false;
                break;
            case '\\':
                ch = escape(str, i);
            default:
                {
                    std::string chars(1, ch);
                    Trigger trigger(chars, kNone);
                    uint new_state = nfa.add_state();
                    nfa.add_transition(prev_end, new_state, trigger);
                    prev_start = prev_end;
                    prev_end = new_state;
                    permit_repeat = true;
                }
                break;
        }
    }
    
    if (!group_end_index.empty())
    {
        // Todo: log error
        return false;
    }

    for (uint i = 0; i < group_end.size(); ++i)
    {
        nfa.set_accept(group_end[i]);
    }
    nfa.set_accept(prev_end);

    return true;
}

char NfaBuilder::escape(const std::string &str, uint &i)
{
    assert(str[i] == '\\');
    return i + 1 < str.size() ? str[++i] : str[i];
}

bool NfaBuilder::parse_char_set(const std::string &str, uint &i, std::string &chars)
{
    bool accepted = false;
    bool selected[256] = {0};

    for (; i < str.size(); ++i)
    {
        char ch = str[i];
        if (ch == ']')
        {
            accepted = true;
            break;
        }

        if (ch == '\\')
        {
            ch = escape(str, i);
        }

        if (i + 1 < str.size() && str[i + 1] == '-'
            && i + 2 < str.size() && str[i + 2] != ']')
        {
            i += 2;
            char end = str[i];
            if (end == '\\')
            {
                end = escape(str, i);
            }
            for (char j = ch; j <= end; ++j)
            {
                selected[static_cast<int>(j)] = true;
            }
        }
        else
        {
            selected[static_cast<int>(ch)] = true;
        }
    }

    if (!accepted)
    {
        // Todo: Log error
        return false;
    }

    chars.clear();
    for (uint j = 0; j < 256; ++j)
    {
        if (selected[j])
        {
            chars.append(1, static_cast<char>(j));
        }
    }

    if (chars.empty())
    {
        // Todo: log error
        return false;
    }
    
    return true;
}

}
